# -*- coding: utf-8 -*-
# created on 2017/3/1
# 

import collections
from sympy.abc import a, b
from sympy import expand
from mathsolver.functions.base import *
from mathsolver.functions.base.base import new_latex


# 求展开式的第r+1项
class EXSRPoly(BaseFunction):
    def solver(self, *args):  # in: comm, r; out: baseValue
        r = args[0].sympify()
        exs = self.search(str('exs'))
        comm = exs.common()
        comm = comm.sympify()
        if isinstance(r, list):
            new_comm = comm.subs({sympify('r'): r[0] - 1})
            self.steps.append(["", "展开式的第%s项为：%s" % (new_latex(r[0] + 1), new_latex(new_comm))])
        else:
            new_comm = comm.subs({sympify('r'): r - 1})
            self.steps.append(["", "展开式的第%s项为：%s" % (new_latex(r + 1), new_latex(new_comm))])
        self.output.append(BaseNumbers(new_comm))
        self.label.add("求展开式的第r+1项")
        return self


# 求系数为有理数的项
class EXSYouLiPolyR(BaseFunction):
    def solver(self, *args):  # in: comm, out: BaseNumber: 个数
        if not args:
            exs = self.search(str('exs'))
            comm = exs.common()
            comm = comm.sympify()
        else:
            raise Exception("to do")
        pow_dict = comm.as_powers_dict()
        nn = -1
        for k, v in pow_dict.items():
            if isinstance(k, binomial):
                nn = k.args[0]
        rs = []
        for r in range(nn + 1):
            new_comm = comm.subs({Symbol("r"): r})
            new_comm_pow_dict = new_comm.as_powers_dict()
            flag = []
            for bb, exp in new_comm_pow_dict.items():
                if exp.is_Integer:
                    flag.append(0)
                else:
                    flag.append(1)
            if sum(flag) == 0:
                rs.append(r + 1)
        self.steps.append(["", "系数为有理数的项分别是第%s项" % ",".join([new_latex(num) for num in rs])])
        self.output.append(BaseNumbers(rs))
        self.label.add("求系数为有理数的项")
        return self


# 求项数
class EXSPolyNum(BaseFunction):
    def solver(self, *args):
        poly_r = args[0].sympify()
        num = len(poly_r)
        self.steps.append(["", "∴共有%s项" % (new_latex(num))])
        self.output.append(BaseNumber(num))
        self.label.add("求展开式的项数")
        return self


# 求常数项
class EXSConstPoly(BaseFunction):
    def solver(self, *args):
        exs = self.search(str('exs'))
        if exs.flag:
            comm = exs.common()
            poly = comm.sympify()
            self.steps.append(["", "依题意, 得"])
            self.steps.append(["", "%s = %s" % (new_latex("T_{r+1}"), new_latex(poly))])
            poly = sympify(poly)
            pow_dict = poly.as_powers_dict()
            eqs = []
            eqs_print = []
            for k, v in pow_dict.items():
                if k in {x, y}:
                    eqs.append(v)
                    eqs_print.append([v, S.Zero])
            if len(eqs_print) == 1:
                self.steps.append(["", "∴%s" % BaseEq(eqs_print[0]).printing()])
            else:
                self.steps.append(["", "∴%s" % BaseEqs(eqs_print).printing()])
            solutions = solve(eqs, sympify('r'), dict=True)
            answer = {}
            key_r = sympify('r')
            values_r = []
            for ans in solutions:
                r_value = ans[key_r]
                r_symbols = r_value.free_symbols
                if r_symbols:
                    values_r.append(ans[key_r])
                else:
                    if r_value.is_Integer:
                        values_r.append(ans[key_r])
            answer[key_r] = values_r
            known_sym, known_solve = zip(*answer.items())
            answer = collections.defaultdict(list)
            for known in zip(*known_solve):
                answer[known_sym].append(tuple(known))
            self.steps.append(["", "解得：%s" % (BaseSymbolValue(answer).printing())])
            if len(values_r) == 1:
                new_poly = poly
                new_poly = new_poly.subs({sympify('r'): values_r[0]})
                self.steps.append(["", "∴常数项为: %s" % (BaseNumber(new_poly).printing())])
                self.output.append(BaseNumber(new_poly))
        else:
            poly = exs.poly
            poly = simplify(poly)
            poly = expand(poly)
            if len(args) == 2:
                sym = args[1].value2pairs()[0]
                poly = poly.subs(sym)
            const = poly.as_coeff_Add()[0]
            self.steps.append(["", "常数项为: %s" % new_latex(const)])
            self.output.append(BaseNumber(const))
        self.label.add("求常数项")
        return self


# 求有理项的r
class EXSPolyExpIntegerR(BaseFunction):
    def solver(self, *args):  # in: comm ， out: BaseNumber项数
        exs = self.search(str('exs'))
        comm_poly = exs.common()
        comm_poly = comm_poly.sympify()
        symbol = args[0].sympify()
        pow_dict = comm_poly.as_powers_dict()
        exp = []
        nn = -1
        for k, v in pow_dict.items():
            if k == symbol:
                exp.append(v)
            elif isinstance(k, binomial):
                nn = k.args[0]
        indices = []
        exp = sum(exp)
        for i in range(nn + 1):
            sub_exp = exp.subs({Symbol("r"): i})
            if sub_exp.is_Integer:
                indices.append(i + 1)
        self.steps.append(
            ["", "%s的幂指数是整数的项是第%s项" % (new_latex(symbol), ",".join([new_latex(num) for num in indices]))])
        self.output.append(BaseNumbers(indices))
        self.label.add("求有理项")
        return self


# 中间项的项数
class ExsMiddleR(BaseFunction):
    def solver(self, *args):  # in: baseSympbol_value({a:,b:,n:})
        self.label.add("求展开式的中间项")
        coeffs = args[0].value2pairs()[0]
        e_a, e_b, e_n = coeffs[a], coeffs[b], int(coeffs[Symbol("n")])
        if e_n & 0x1 == 1:  # 奇数
            nums = [e_n / 2 + 1, e_n / 2 + 2]
            self.steps.append(["", "中间项为第%s,%s项" % (new_latex(nums[0]), new_latex(nums[1]))])
        else:
            nums = [e_n / 2 + 1]
            self.steps.append(["", "中间项为第%s项" % new_latex(nums[0])])
        self.output.append(BaseNumbers(nums))
        return self


# 求是包含某poly的项
# 设i为虚数单位,则(x+i)^{6}的展开式中含x^{4}的项为 ( )
class EXSIncludePolyIsWhat(BaseFunction):
    def solver(self, *args):  # 输入comm, poly,输出BaseSymbolValue, r
        exs = self.search(str('exs'))
        base_poly = exs.poly
        base_poly = simplify(base_poly)
        poly_ed = base_poly.expand()
        poly = args[0].sympify()
        if not poly_ed.is_Add:
            base_poly, poly = poly, base_poly
            poly_ed = base_poly.expand()
        coef = poly_ed.coeff(poly)
        xiang = coef * poly
        self.steps.append(["", "∴含%s的项为: %s" % (new_latex(poly), new_latex(xiang))])
        self.output.append(BasePoly(xiang))
        self.label.add("求展开式包含某单项式的项")
        return self


if __name__ == '__main__':
    pass
